/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "bsr_common.h"

namespace YHAMG
{

void VecBlockFill(int block_size, double a, double *x)
{
	for (int i = 0; i < block_size; ++i)
		x[i] = a;
}

void VecBlockScale(int block_size, double a, double *x)
{
	for (int i = 0; i < block_size; ++i)
		x[i] *= a;
}

void VecBlockCopy(int block_size, const double *x, double *y)
{
	for (int i = 0; i < block_size; ++i)
		y[i] = x[i];
}

void VecBlockAdd(int block_size, const double *x, double *y)
{
	for (int i = 0; i < block_size; ++i)
		y[i] += x[i];
}

void VecBlockSub(int block_size, const double *x, double *y)
{
	for (int i = 0; i < block_size; ++i)
		y[i] -= x[i];
}

void VecBlockScaleAdd(int block_size, double a, const double *x, double *y)
{
	for (int i = 0; i < block_size; ++i)
		y[i] += a * x[i];
}

void BSRBlockFill(int block_size, double a, double *A)
{
	int block_nnz = block_size * block_size;
	for (int i = 0; i < block_nnz; ++i)
		A[i] = a;
}

void BSRBlockScale(int block_size, double a, double *A)
{
	int block_nnz = block_size * block_size;
	for (int i = 0; i < block_nnz; ++i)
		A[i] *= a;
}

void BSRBlockCopy(int block_size, const double *A, double *B)
{
	int block_nnz = block_size * block_size;
	for (int i = 0; i < block_nnz; ++i)
		B[i] = A[i];
}

void BSRBlockAdd(int block_size, const double *A, double *B)
{
	int block_nnz = block_size * block_size;
	for (int i = 0; i < block_nnz; ++i)
		B[i] += A[i];
}

void BSRBlockSub(int block_size, const double *A, double *B)
{
	int block_nnz = block_size * block_size;
	for (int i = 0; i < block_nnz; ++i)
		B[i] -= A[i];
}

void BSRBlockScaleAdd(int block_size, double a, const double *A, double *B)
{
	int block_nnz = block_size * block_size;
	for (int i = 0; i < block_nnz; ++i)
		B[i] += a * A[i];
}

void BSRBlockMatVec(int block_size, const double *A, const double *x, double* y)
{
	for (int i = 0; i < block_size; ++i)
	{
		double temp = 0.0;
		for (int j = 0; j < block_size; ++j)
			temp += A[i + j * block_size] * x[j];
		y[i] = temp;
	}
}

void BSRBlockMatVecAdd(int block_size, const double *A, const double *x, double* y)
{
	for (int i = 0; i < block_size; ++i)
	{
		double temp = y[i];
		for (int j = 0; j < block_size; ++j)
			temp += A[i + j * block_size] * x[j];
		y[i] = temp;
	}
}

void BSRBlockMatVecSub(int block_size, const double *A, const double *x, double* y)
{
	for (int i = 0; i < block_size; ++i)
	{
		double temp = y[i];
		for (int j = 0; j < block_size; ++j)
			temp -= A[i + j * block_size] * x[j];
		y[i] = temp;
	}
}

void BSRBlockTranspose(int block_size, const double *A, double *B)
{
	for (int i = 0; i < block_size; ++i)
		for (int j = 0; j < block_size; ++j)
			B[i + j * block_size] = A[i * block_size + j];
}

void BSRBlockMult(int block_size, const double *A, const double *B, double *C)
{
	for (int j = 0; j < block_size; ++j)
	{
		for (int i = 0; i < block_size; ++i)
		{
			double temp = 0.0;
			for (int k = 0; k < block_size; ++k)
				temp += A[i + k * block_size] * B[k + j * block_size];
			C[i + j * block_size] = temp;
		}
	}
}

void BSRBlockMultAdd(int block_size, const double *A, const double *B, double *C)
{
	for (int j = 0; j < block_size; ++j)
	{
		for (int i = 0; i < block_size; ++i)
		{
			double temp = C[i + j * block_size];
			for (int k = 0; k < block_size; ++k)
				temp += A[i + k * block_size] * B[k + j * block_size];
			C[i + j * block_size] = temp;
		}
	}
}

void BSRBlockMultSub(int block_size, const double *A, const double *B, double *C)
{
	for (int j = 0; j < block_size; ++j)
	{
		for (int i = 0; i < block_size; ++i)
		{
			double temp = C[i + j * block_size];
			for (int k = 0; k < block_size; ++k)
				temp -= A[i + k * block_size] * B[k + j * block_size];
			C[i + j * block_size] = temp;
		}
	}
}

void BSRBlockLUFactorize(int block_size, double *LU)
{
	for (int k = 0; k < block_size; ++k)
	{
		double temp = 1.0 / LU[k * block_size + k];
		LU[k * block_size + k] = temp;
		for (int i = k + 1; i < block_size; ++i)
		{
			LU[i + k * block_size] = LU[i + k * block_size] * temp;
			for (int j = k + 1; j < block_size; ++j)
				LU[i + j * block_size] -= LU[i + k * block_size] * LU[k + j * block_size];
		}
	}
}

void BSRBlockLUVecSolve(int block_size, const double *LU, double* x)
{
	for (int i = 0; i < block_size; ++i)
	{
		double temp = x[i];
		for (int k = 0; k < i; ++k)
			temp -= x[k] * LU[i + k * block_size];
		x[i] = temp;
	}

	for (int i = block_size - 1; i >= 0; --i)
	{
		double temp = x[i];
		for (int k = block_size - 1; k > i; --k)
			temp -= x[k] * LU[i + k * block_size];
		x[i] = temp * LU[i + i * block_size];
	}
}

void BSRBlockLUMatSolve(int block_size, const double *LU, double* X)
{
	for (int j = 0; j < block_size; ++j)
	{
		for (int i = 0; i < block_size; ++i)
		{
			double temp = X[i + j * block_size];
			for (int k = 0; k < i; ++k)
				temp -= X[k + j * block_size] * LU[i + k * block_size];
			X[i + j * block_size] = temp;
		}
	}
	
	for (int j = block_size - 1; j >= 0; --j)
	{
		for (int i = block_size - 1; i >= 0; --i)
		{
			double temp = X[i + j * block_size];
			for (int k = block_size - 1; k > i; --k)
				temp -= X[k + j * block_size] * LU[i + k * block_size];
			X[i + j * block_size] = temp * LU[i + i * block_size];
		}
	}
}

void BSRBlockMatLUSolve(int block_size, const double *LU, double* X)
{
	for (int j = 0; j < block_size; ++j)
	{
		for (int i = 0; i < block_size; ++i)
		{
			double temp = X[j + i * block_size];
			for (int k = 0; k < i; ++k)
				temp -= X[j + k * block_size] * LU[k + i * block_size];
			X[j + i * block_size] = temp * LU[i + i * block_size];
		}
	}

	for (int j = block_size - 1; j >= 0; --j)
	{
		for (int i = block_size - 1; i >= 0; --i)
		{
			double temp = X[j + i * block_size];
			for (int k = block_size - 1; k > i; --k)
				temp -= X[j + k * block_size] * LU[k + i * block_size];
			X[j + i * block_size] = temp;
		}
	}
}

double BSRBlockSum(int block_size, const double* A)
{
	int block_nnz = block_size * block_size;
	double result = 0.0;
	for (int i = 0; i < block_nnz; ++i)
		result += A[i];
	return result;
}

}